import { NextRequest, NextResponse } from "next/server";
import { Message as VercelChatMessage, StreamingTextResponse } from "ai";
import { AgentExecutor, createOpenAIFunctionsAgent } from "langchain/agents";
import { ChatOpenAI, OpenAIEmbeddings } from "@langchain/openai";
import { Calculator } from "@langchain/community/tools/calculator";
import { SerpAPI } from "@langchain/community/tools/serpapi";
import { AIMessage, ChatMessage, HumanMessage } from "@langchain/core/messages";
import * as hub from "langchain/hub";;
import { z } from "zod";
import { DynamicTool, DynamicStructuredTool } from "@langchain/core/tools";
import { createRetrieverTool } from "langchain/tools/retriever";

import { ChatPromptTemplate, /* MessagesPlaceholder */ } from "@langchain/core/prompts";
import { PineconeStore } from "@langchain/pinecone";
import { PINECONE_INDEX_NAME, PINECONE_NAME_SPACE } from "@/config/pinecone";
import { pinecone } from "@/lib/utils/pinecone-client";

export const runtime = "edge";

const convertVercelMessageToLangChainMessage = (message: VercelChatMessage) => {
  if (message.role === "user") {
    return new HumanMessage(message.content);
  } else if (message.role === "assistant") {
    return new AIMessage(message.content);
  } else {
    return new ChatMessage(message.content, message.role);
  }
};

// const prompt = ChatPromptTemplate.fromMessages([
//     ["system", "You are a helpful assistant"],
//     ["placeholder", "{chat_history}"],
//     ["human", "{input}"],
//     new MessagesPlaceholder("agent_scratchpad"),
// ]);

/**
 * https://js.langchain.com/docs/modules/agents/agent_types/openai_functions_agent
 */
export async function POST(req: NextRequest) {
  try {
    const body = await req.json();
    const messages = (body.messages ?? []).filter(
      (message: VercelChatMessage) =>
        message.role === "user" || message.role === "assistant",
    );
    const returnIntermediateSteps = body.show_intermediate_steps;
    const previousMessages = messages
      .slice(0, -1)
      .map(convertVercelMessageToLangChainMessage);
    const currentMessageContent = messages[messages.length - 1].content;

    const prompt = await hub.pull<ChatPromptTemplate>("hwchase17/openai-functions-agent");
    const llm = new ChatOpenAI({
      model: "gpt-4-turbo",
      temperature: 0,
      // streaming: true,
    });

    // 读取向量库
    const vectorStore = await PineconeStore.fromExistingIndex(
      new OpenAIEmbeddings(),
      {
        pineconeIndex: pinecone.Index(PINECONE_INDEX_NAME),
        namespace: PINECONE_NAME_SPACE,
        textKey: 'text'
      }
    );
    const retrieverTool = createRetrieverTool(vectorStore.asRetriever(), {
      name: "search_latest_knowledge",
      description: "搜索知识库获取最新信息。",
    });

    // 定义工具
    // https://serpapi.com/
    const tools = [
      // retrieverTool, // 知识库
      // new Calculator(), // 计算
      // new SerpAPI(), // 搜索引擎
      new DynamicTool({
        name: "current_date",
        description: "获取当前日期。input 应该是个空字符串",
        func: async () => {
          const date = new Date();
          return date.toLocaleDateString();
        },
      }),
      new DynamicStructuredTool({
        name: "random_number",
        description: "在两个输入数字之间生成一个随机数",
        schema: z.object({
          low: z.number().describe("生成随机数的下限"),
          high: z.number().describe("生成随机数的上限"),
        }),
        func: async ({ low, high }) =>
          (Math.random() * (high - low) + low).toString(), // Outputs still must be strings
      }),
    ];
    const agent = await createOpenAIFunctionsAgent({
      llm,
      tools,
      prompt,
    });

    const agentExecutor = new AgentExecutor({
      agent,
      tools,
      returnIntermediateSteps,
      verbose: false,
    });
    if (!returnIntermediateSteps) {
      const orginStream = await agentExecutor.stream({
        input: currentMessageContent,
        chat_history: previousMessages,
      });
      const textEncoder = new TextEncoder();
      const stream = new ReadableStream({
        async start(controller) {
          for await (const chunk of orginStream) {
            if (chunk.output) {
              controller.enqueue(textEncoder.encode(chunk.output));
            }
          }
          controller.close();
        },
      });
      return new StreamingTextResponse(stream);
    }
    const result = await agentExecutor.invoke({
      input: currentMessageContent,
      chat_history: previousMessages,
    });
    return NextResponse.json({ output: result.output, intermediate_steps: result.intermediateSteps });
  } catch (e: any) {
    return NextResponse.json({ error: e.message }, { status: e.status ?? 500 });
  }
}
